import pymongo


class MongoOperator:
    def __init__(self, host, port):
        self.host = host
        self.port = port
        self.mongo_client = pymongo.MongoClient(host=host, port=port)

    # 检查数据库是否存在
    def check_db_exist(self, db_name):
        dblist = self.mongo_client.list_database_names()
        if db_name in dblist:
            return True
        else:
            print("数据库【" + db_name + "】不存在！")
            return False

    # 检查数据集是否存在
    def check_collection_exist(self, db_name, collection_name):
        if self.check_db_exist(db_name):
            db = self.mongo_client[db_name]
            if collection_name in db.list_collection_names():
                return True
            else:
                print("数据表【" + collection_name + "】不存在！")
                return False
        else:
            return False

    def update_one(self, db_name, collection_name, item):
        try:
            f = {"_id": item["_id"]}
            kv = {}
            for key in item.keys():
                if "_id" == key:
                    continue
                kv[key] = item[key]
            update = {"$set": kv}
            db = self.mongo_client[db_name]
            db[collection_name].update_one(filter=f, update=update)
            return True
        except Exception as e:
            # 如果不是因为主键重复导致的错误，输出错误信息
            if "duplicate key error index" not in str(e):
                print("数据更新错误********************************************************")
                print(item)
                print(e)
                print(kv)
                print(f)
                return False
            return True

    def insert_or_update_one(self, db_name, collection_name, item):
        try:
            db = self.mongo_client[db_name]
            db[collection_name].insert_one(item)
            return True
        except Exception as e:
            # 如果不是因为主键重复导致的错误，输出错误信息
            if "duplicate key error" not in str(e):
                print("数据插入或更新错误********************************************************")
                print(item)
                print(e)
                raise e
            else:
                self.update_one(db_name, collection_name, item)

    def insert(self, db_name, collection_name, items):
        if isinstance(items, list):
            for item in items:
                # 即使是多条一次也只能插入一条，插入多条如果某条有重复则全部插入失败
                self.insert_or_update_one(db_name, collection_name, item)
            # self.insert_many(db_name, collection_name, items)
        elif isinstance(items, dict):
            self.insert_or_update_one(db_name, collection_name, items)
        else:
            print("数据类型不符合要求！")
            print(db_name, collection_name, items)
